import pymysql
import redis
from pymysql import MySQLError
import time, datetime
"""
数据库更新代码：
此部分主要是用来更新权限数据库记录。
更新数据库时，要考虑redis和MySQL的数据一致性问题。所以在更新数据时，先查询redis中是否有此记录，若有，
则先删除redis记录，再插入数据到redis，然后将数据插入MySQL中，插入成功后，重新将数据再次更新到redis。
若redis中无记录，则直接更新MySQL中的记录。
"""


class UpdateDB:
    def __init__(self):
        # 连接MySQL数据库
        try:
            self.conn = pymysql.connect(host='1.1.1.1', port=3306,
                                        user='root', password='111111',
                                        database='auth_data_db', charset='utf8')
        except Exception as error:
            print('连接MySQL出现问题！')
            print('失败原因：', error)
            exit()

        try:
            # 建立redis连接池
            self.conn_pool = redis.ConnectionPool(host='1.1.1.1', port=6379, db=0, decode_responses=True,
                                                  password='111111')
            # 客户端0连接数据库
            self.r0 = redis.StrictRedis(connection_pool=self.conn_pool)
        except Exception as error:
            print('连接redis出现问题！')
            print('失败原因：', error)
            exit()

    """
    更新数据的操作，为了避免更新MySQL后，redis没更新的这一段空挡时间的查询，所以先更新redis，
    再更新MySQL，然后MySQL成功提交后，再次对redis进行重新更新
    """

    def post_data(self):
        # 插入数据
        print('输入更新数据：')
        src_ip = input('src_ip:')
        dst_ip = input('dst_ip:')
        dst_port = input('dst_port:')
        auth = input('auth:')
        # redis  string类型表的key
        key = src_ip + ':' + dst_ip + ':' + dst_port

        # 先查询redis数据库是否存在数据,如果存在数据则更新redis，再更新MySQL，若不存在则去MySQL中更新,提交成功再次更新redis
        result = self.r0.get(key)
        # reids存在数据，则需要对数据进行更新，即先清除再写入; 写入redis后，再将数据写入MySQL
        if result:
            # 清除数据
            self.r0.delete(key)
            self.r0.set(key, auth)
            self.r0.expire(key, 600)  # 设置过期时间

            # 先查询MySQL表内是否存在记录，如不存在记录，则插入记录；如存在记录，则修改记录
            count = 0
            with self.conn.cursor() as find_cursor:
                try:
                    count = find_cursor.execute(
                        'select count(rno) from tb_as where src_ip=%s and dst_ip=%s and dst_port=%s',
                        (src_ip, dst_ip, int(dst_port),)
                    )
                except MySQLError as error:
                    print(error)

            with self.conn.cursor() as cursor:
                try:
                    if count == 0:
                        # 插入SQL语句，result为返回的结果
                        res_info = cursor.execute(
                            'insert into tb_student values (%s, %s, %s, %s)',
                            (src_ip, dst_ip, int(dst_port), auth,)
                        )
                    else:
                        # 更新权限数据
                        res_info = cursor.execute(
                            'UPDATE tb_as SET acc_auth=%s WHERE src_ip=%s AND dst_ip=%s AND dst_port=%s',
                            (auth, src_ip, dst_ip, int(dst_port),)
                        )
                    # 成功插入后需要提交才能同步在数据库中
                    if isinstance(res_info, int):
                        print('数据更新成功')
                        self.conn.commit()
                        # 再次更新redis
                        self.r0.set(key, auth)
                        self.r0.expire(key, 600)  # 设置过期时间
                except MySQLError as error:
                    # 如果MySQL提交不成功，清除redis数据
                    self.r0.delete(key)
                    print(error)
                    self.conn.rollback()
                finally:
                    # 操作执行完成后，需要关闭连接
                    self.conn.close()
        else:
            # 先查询MySQL表内是否存在记录，如不存在记录，则插入记录；如存在记录，则修改记录
            count = 0
            with self.conn.cursor() as find_cursor:
                try:
                    count = find_cursor.execute(
                        'select count(rno) from tb_as where src_ip=%s and dst_ip=%s and dst_port=%s',
                        (src_ip, dst_ip, int(dst_port),)
                    )
                except MySQLError as error:
                    print(error)

            with self.conn.cursor() as cursor:
                try:
                    if count == 0:
                        # 插入SQL语句，result为返回的结果
                        res_info = cursor.execute(
                            'insert into tb_student values (%s, %s, %s, %s)',
                            (src_ip, dst_ip, int(dst_port), auth,)
                        )
                    else:
                        # 更新权限数据
                        res_info = cursor.execute(
                            'UPDATE tb_as SET acc_auth=%s WHERE src_ip=%s AND dst_ip=%s AND dst_port=%s',
                            (auth, src_ip, dst_ip, int(dst_port),)
                        )
                    # 成功插入后需要提交才能同步在数据库中
                    if isinstance(res_info, int):
                        print('数据更新成功')
                        self.conn.commit()
                except MySQLError as error:
                    print(error)
                    self.conn.rollback()
                finally:
                    # 操作执行完成后，需要关闭连接
                    self.conn.close()


if __name__ == '__main__':
    dbs = UpdateDB()
    dbs.post_data()
